// fitsconverter.cpp
#include "fitsconverter.h"
#include "fitsio.h"
#include <fstream>
#include <stdexcept>
#include <vector>
#include <cstring>
#include <algorithm>
#include <memory>
#include <sys/stat.h>

// 读取二进制文件
std::vector<uint16_t> FitsConverter::readBinaryFile(const std::string& filename) {
    std::ifstream file(filename, std::ios::binary | std::ios::ate);
    if (!file) {
        throw std::runtime_error("无法打开文件: " + filename);
    }

    std::streamsize size = file.tellg();
    file.seekg(0, std::ios::beg);

    if (size % sizeof(uint16_t) != 0) {
        throw std::runtime_error("文件大小不是uint16_t的整数倍");
    }

    std::vector<uint16_t> buffer(size / sizeof(uint16_t));
    if (!file.read(reinterpret_cast<char*>(buffer.data()), size)) {
        throw std::runtime_error("读取文件失败: " + filename);
    }

    // 确保小端序读取（如果Python版本是小端序）
    // 根据您的Python脚本使用的字节序决定是否需要交换
#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
    for (auto& val : buffer) {
        val = __builtin_bswap16(val);
    }
#endif

    return buffer;
}

// 分割数据到列
std::vector<std::vector<uint16_t>> FitsConverter::splitDataIntoColumns(
    const std::vector<uint16_t>& data, int numColumns) {

    if (data.size() % numColumns != 0) {
        throw std::runtime_error("数据长度必须是列数的整数倍");
    }

    size_t elementsPerColumn = data.size() / numColumns;
    std::vector<std::vector<uint16_t>> columns(numColumns, std::vector<uint16_t>(elementsPerColumn));

    // 更高效的列分割方式，确保顺序与Python一致
    for (size_t i = 0; i < elementsPerColumn; ++i) {
        for (int j = 0; j < numColumns; ++j) {
            columns[j][i] = data[i * numColumns + j];
        }
    }

    return columns;
}

// 重塑列数据 (4700x1250)
std::vector<std::vector<uint16_t>> FitsConverter::reshapeColumns(
    const std::vector<std::vector<uint16_t>>& columns, int rows, int cols) {

    std::vector<std::vector<uint16_t>> reshapedColumns;
    reshapedColumns.reserve(columns.size());

    for (const auto& col : columns) {
        if (col.size() != static_cast<size_t>(rows * cols)) {
            throw std::runtime_error("列数据大小与重塑尺寸不匹配");
        }

        std::vector<uint16_t> reshaped(rows * cols);
        for (int i = 0; i < rows; ++i) {
            for (int j = 0; j < cols; ++j) {
                reshaped[i * cols + j] = col[j * rows + i]; // 转置重塑
            }
        }
        reshapedColumns.push_back(reshaped);
    }

    return reshapedColumns;
}

// 拼接列数据 (2行8列)，完全按照Python版本的逻辑
std::vector<uint16_t> FitsConverter::concatenateColumns(
    const std::vector<std::vector<uint16_t>>& columns) {

    const int numRows = 2;
    const int numColumnsPerRow = 8;
    const int colHeight = 4700;
    const int colWidth = 1250;

    if (columns.size() != 16) {
        throw std::runtime_error("需要16列数据");
    }

    // 验证每列数据大小
    for (const auto& col : columns) {
        if (col.size() != colHeight * colWidth) {
            throw std::runtime_error("列数据大小不正确");
        }
    }

    // 计算最终图像尺寸
    const int totalWidth = colWidth * numColumnsPerRow;
    const int totalHeight = colHeight * numRows;
    std::vector<uint16_t> combinedImage(totalWidth * totalHeight, 0);

    // 定义与Python版本完全相同的通道重排顺序
    const std::vector<int> reorder_lower = {2, 3, 0, 1, 6, 7, 4, 5};  // 对应Python的[3,4,1,2,7,8,5,6]
    const std::vector<int> reorder_upper = {5, 4, 7, 6, 1, 0, 3, 2};  // 对应Python的[c,d,a,b,g,h,e,f]

    for (int row = 0; row < numRows; ++row) {
        for (int col = 0; col < numColumnsPerRow; ++col) {
            // 计算原始通道索引
            int origChan;
            if (row == 0) {
                origChan = reorder_lower[col];
            } else {
                origChan = 8 + reorder_upper[col]; // 上排通道
            }

            if (origChan < 0 || origChan >= 16) {
                throw std::runtime_error("无效的通道索引");
            }

            const auto& column = columns[origChan];

            // 计算目标位置
            int destX = col * colWidth;
            int destY = row * colHeight;

            // 复制数据，同时应用翻转
            for (int y = 0; y < colHeight; ++y) {
                for (int x = 0; x < colWidth; ++x) {
                    // 计算源位置，考虑翻转
                    int srcY = y;
                    int srcX = x;

                    // 底行需要垂直翻转
                    if (row == 1) {
                        srcY = colHeight - 1 - y;
                    }

                    // 右半图像水平翻转 (通道4-7和12-15)
                    if (origChan % 8 >= 4) {
                        srcX = colWidth - 1 - x;
                    }

                    // 计算源位置索引
                    int srcPos = srcY * colWidth + srcX;

                    // 确保不越界
                    if (srcPos >= 0 && srcPos < static_cast<int>(column.size())) {
                        int destPos = (destY + y) * totalWidth + (destX + x);
                        if (destPos >= 0 && destPos < static_cast<int>(combinedImage.size())) {
                            combinedImage[destPos] = column[srcPos];
                        }
                    }
                }
            }
        }
    }

    return combinedImage;
}

// 写入FITS文件
bool FitsConverter::writeFitsFile(const std::string& filePath,
                                  const std::vector<uint16_t>& imageData,
                                  int width, int height,
                                  const std::vector<std::pair<std::string, std::string>>& keywords) {
    fitsfile* fitsFile = nullptr;
    int status = 0;

    // 删除已存在的文件
    remove(filePath.c_str());

    // 创建新FITS文件
    fits_create_file(&fitsFile, ("!" + filePath).c_str(), &status);
    if (status) {
        char errText[FLEN_STATUS];
        fits_get_errstatus(status, errText);
        throw std::runtime_error("无法创建FITS文件: " + std::string(errText));
    }

    // 创建图像HDU
    long naxes[2] = {width, height};
    fits_create_img(fitsFile, USHORT_IMG, 2, naxes, &status);
    if (status) {
        fits_report_error(stderr, status);
        throw std::runtime_error("无法创建图像HDU");
    }

    // 添加关键字
    for (const auto& kv : keywords) {
        std::string key = kv.first.substr(0, 8); // FITS关键字最多8字符
        std::string value = kv.second;
        std::string comment = "";

        fits_update_key(fitsFile, TSTRING, key.c_str(), (void*)value.c_str(), comment.c_str(), &status);
        if (status) {
            fits_report_error(stderr, status);
            throw std::runtime_error("无法写入关键字");
        }
    }

    // 写入图像数据
    long fpixel[2] = {1, 1};
    fits_write_pix(fitsFile, TUSHORT, fpixel, imageData.size(), (void*)imageData.data(), &status);
    if (status) {
        fits_report_error(stderr, status);
        throw std::runtime_error("无法写入图像数据");
    }

    // 关闭FITS文件
    fits_close_file(fitsFile, &status);
    if (status) {
        fits_report_error(stderr, status);
        throw std::runtime_error("无法关闭FITS文件");
    }

    return true;
}

// 主处理函数（匹配Python脚本）
bool FitsConverter::processBinToFits(const std::string& binPath,
                                     const std::string& outputDir,
                                     const std::string& baseFilename,
                                     const std::vector<std::pair<std::string, std::string>>& keywords) {
    try {
        // 确保输出目录存在
        if (mkdir(outputDir.c_str(), 0755) != 0 && errno != EEXIST) {
            throw std::runtime_error("无法创建输出目录: " + outputDir);
        }

        // 1. 读取二进制文件
        auto data = readBinaryFile(binPath);

        // 2. 分割数据为16列
        auto columns = splitDataIntoColumns(data, 16);

        // 3. 重塑每列为4700x1250
        auto reshapedColumns = reshapeColumns(columns, 4700, 1250);

        // 4. 拼接为2行8列的图像，完全按照Python版本的逻辑
        auto combinedImage = concatenateColumns(reshapedColumns);

        // 5. 保存为FITS文件
        std::string outputFile = outputDir + "/" + baseFilename + ".fits";
        bool success = writeFitsFile(outputFile, combinedImage, 1250 * 8, 4700 * 2, keywords);

        return success;
    } catch (const std::exception& e) {
        throw std::runtime_error("BIN转FITS失败: " + std::string(e.what()));
    }
}
